/** @file
  BDS routine to process capsules

 @copyright
  INTEL CONFIDENTIAL
  Copyright 2013 - 2017 Intel Corporation.

  The source code contained or described herein and all documents related to the
  source code ("Material") are owned by Intel Corporation or its suppliers or
  licensors. Title to the Material remains with Intel Corporation or its suppliers
  and licensors. The Material may contain trade secrets and proprietary and
  confidential information of Intel Corporation and its suppliers and licensors,
  and is protected by worldwide copyright and trade secret laws and treaty
  provisions. No part of the Material may be used, copied, reproduced, modified,
  published, uploaded, posted, transmitted, distributed, or disclosed in any way
  without Intel's prior express written permission.

  No license under any patent, copyright, trade secret or other intellectual
  property right is granted to or conferred upon you by disclosure or delivery
  of the Materials, either expressly, by implication, inducement, estoppel or
  otherwise. Any license under such intellectual property rights must be
  express and approved by Intel in writing.

  Unless otherwise agreed by Intel in writing, you may not remove or alter
  this notice or any other notice embedded in Materials by Intel or
  Intel's suppliers or licensors in any way.

  This file contains a 'Sample Driver' and is licensed as such under the terms
  of your license agreement with Intel or your vendor. This file may be modified
  by the user, subject to the additional terms of the license agreement.

@par Specification Reference:
**/

#include "BdsPlatform.h"
#include <Protocol/CapsuleStatusProtocol.h>
#include <Guid/CapsuleProcessingResult.h>
#include <Library/EcMiscLib.h>
#include <Library/Tpm2CommandLib.h>
#include <Guid/TpmInstance.h>
#include <PchResetPlatformSpecific.h>

#pragma pack(1)
typedef struct {
  UINT32   CapsuleArrayNumber;
  VOID*    CapsulePtr[1];
} CRASHDUMP_CAPSULE_TABLE;
#pragma pack()

#define MAX_CAPSULE_SUPPORTED 10 // The Max number of Capsules BIOS supported

/**

  Check if there is any existing "CapsuleXXXX" Variables.
  If it exist, clear it.

  @retval EFI_SUCCESS             All "CapsuleXXXX" Variables are cleared.

**/
EFI_STATUS
ClearCapsuleStatusVariable (
  VOID
  )
{
  EFI_STATUS Status;
  UINT32 Index;
  UINTN Size;
  CHAR16 CapsuleVariableName [sizeof ("Capsule####")];
  EFI_CAPSULE_PROCESSING_RESULT CapsuleProcessingResult;

  for (Index = 0; Index <= MAX_CAPSULE_SUPPORTED; Index++ ) {
    UnicodeSPrint (CapsuleVariableName, sizeof (CapsuleVariableName), L"Capsule%04x", Index);

    Size = sizeof (EFI_CAPSULE_PROCESSING_RESULT);
    Status = gRT->GetVariable (
                    CapsuleVariableName,
                    &gEfiCapsuleReportGuid,
                    NULL,
                    &Size,
                    &CapsuleProcessingResult);
    if (Status == EFI_SUCCESS) {
      Status = gRT->SetVariable (
                      CapsuleVariableName,
                      &gEfiCapsuleReportGuid,
                      0,
                      0,
                      NULL
                      );
    } else {
      break;
    }
  }
  return EFI_SUCCESS;
}

/**
  Create the Variable "CapsuleXXXX" as defined in UEFI spec 2.4 section 7.5.6.

  @param[in] MostRecentProcessedCapsuleCount  The number count of the processed capsule.
  @param[in] CapsuleHeader                    Pointer to the Capsule Header.
  @param[in] CapsuleStatus                    The process result of this Capsule.

  @retval EFI_SUCCESS                         Variable is created successfully.
**/
EFI_STATUS
CreateCapsuleStatusVariable (
  IN UINT32 MostRecentProcessedCapsuleCount,
  IN EFI_CAPSULE_HEADER *CapsuleHeader,
  IN EFI_STATUS CapsuleStatus
  )
{
  EFI_STATUS Status;
  CHAR16 CapsuleVariableName[sizeof ("Capsule####")];
  EFI_CAPSULE_PROCESSING_RESULT CapsuleProcessingResult;

  if (MostRecentProcessedCapsuleCount > MAX_CAPSULE_SUPPORTED) {
    return EFI_OUT_OF_RESOURCES;
  }

  CapsuleProcessingResult.VariableTotalSize = sizeof (CapsuleProcessingResult);
  CopyMem (&CapsuleProcessingResult.CapsuleGuid, &CapsuleHeader->CapsuleGuid, sizeof (EFI_GUID));
  gRT->GetTime ((EFI_TIME *)&CapsuleProcessingResult.CapsuleProcessed.Year, NULL);
  CapsuleProcessingResult.CapsuleStatus = CapsuleStatus;

  UnicodeSPrint (CapsuleVariableName, sizeof (CapsuleVariableName), L"Capsule%04x", MostRecentProcessedCapsuleCount);
  Status = gRT->SetVariable (
                  CapsuleVariableName,
                  &gEfiCapsuleReportGuid,
                  EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                  sizeof (CapsuleProcessingResult),
                  &CapsuleProcessingResult
                  );

  return Status;
}

/**
  Function gets the file information from an open file descriptor, and stores it
  in a buffer allocated from pool.

  @param[in] FHand                A file handle

  @retval UINTN                  File Size
**/
UINTN
GetFileLength (
  IN EFI_FILE_HANDLE              FHand
  )
{
  EFI_STATUS                      Status;
  EFI_FILE_INFO                   *Buffer;
  UINTN                           BufferSize;

  ///
  /// Initialize for GrowBuffer loop
  ///
  BufferSize  = SIZE_OF_EFI_FILE_INFO + 200;

  ///
  /// Call the real function
  ///
  do {
    Buffer = AllocateZeroPool (BufferSize);
    Status = FHand->GetInfo (
                      FHand,
                      &gEfiFileInfoGuid,
                      &BufferSize,
                      Buffer
                      );
    DEBUG ((DEBUG_INFO, "GetInfo %r", Status));
    DEBUG ((DEBUG_INFO, " FileSize %lX (%ld) Buffer%lX\n", BufferSize, BufferSize, Buffer));
    if (!EFI_ERROR (Status)) {
      break;
    }
    if (Status != EFI_BUFFER_TOO_SMALL) {
      FreePool (Buffer);
      break;
    }
    FreePool (Buffer);
  } while (TRUE);

  if (Buffer) {
    BufferSize = Buffer->FileSize;
    FreePool (Buffer);
  }  else {
    BufferSize = 0;
  }
  return BufferSize;
}

EFI_STATUS
AutomaticFirmwareUpdateHandler (
  CHAR16 *FileName
  )
{
  EFI_STATUS                          Status ;
  EFI_HANDLE                          *SimpleFileSystemHandles;
  UINTN                               NumberSimpleFileSystemHandles;
  UINTN                               Index;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL     *Fs;
  EFI_FILE_HANDLE                     Root;
  EFI_FILE_HANDLE                     FileHandle;
  UINTN                               BufferSize;
  VOID                                *Buffer;
  EFI_CAPSULE_HEADER                  *CapsuleHeaderArray[2];
  UINT64                              MaxCapsuleSize;
  EFI_RESET_TYPE                      ResetType;
  EFI_CAPSULE_BLOCK_DESCRIPTOR        *BlockDescriptors;

  DEBUG ((DEBUG_INFO, "AutomaticFirmwareUpdateHandler entry\n"));
  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiSimpleFileSystemProtocolGuid,
                  NULL,
                  &NumberSimpleFileSystemHandles,
                  &SimpleFileSystemHandles
                  );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  for (Index = 0;Index < NumberSimpleFileSystemHandles;Index++) {
    Status = gBS->HandleProtocol (SimpleFileSystemHandles[Index], &gEfiSimpleFileSystemProtocolGuid, (VOID **)&Fs);
    if (EFI_ERROR (Status)) {
      continue;
    }

    Status = Fs->OpenVolume (Fs, &Root);
    if (EFI_ERROR (Status)) {
      continue;
    }

    FileHandle = NULL;
    Status = Root->Open (Root, &FileHandle, FileName, EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE, 0);
    DEBUG ((DEBUG_INFO, "Find %S %r\n", FileName, Status));
    if (EFI_ERROR (Status)) {
      Root->Close (Root);
      continue;
    }
    BufferSize = GetFileLength (FileHandle);
    if (BufferSize == 0) {
      FileHandle->Close (FileHandle);
      DEBUG ((DEBUG_ERROR, "FileSize is 0\n"));
      Root->Close (Root);
      continue;
    }
    Buffer = AllocateZeroPool (BufferSize);
    DEBUG ((DEBUG_INFO, "BufferSize:%lX Buffer:%lX\n", BufferSize, Buffer));
    if (Buffer == NULL) {
      DEBUG ((DEBUG_ERROR, "AllocateZeroPool:Fail\n"));
      FileHandle->Close (FileHandle);
      Root->Close (Root);
      continue;
    }
    Status = FileHandle->Read (FileHandle, &BufferSize, Buffer);
    DEBUG ((DEBUG_ERROR, "Read %S %r\n", FileName, Status));
    DEBUG ((DEBUG_INFO, "BufferSize:%lX (%ld) Buffer:%lX\n", BufferSize, BufferSize, Buffer));
    FileHandle->Delete (FileHandle);
    Root->Close (Root);
    //
    // Inquire platform capability of UpdateCapsule.
    //
    CapsuleHeaderArray[0] = (EFI_CAPSULE_HEADER *) Buffer;
    CapsuleHeaderArray[1] = NULL;
    Status = gRT->QueryCapsuleCapabilities(CapsuleHeaderArray,1,&MaxCapsuleSize,&ResetType);
    DEBUG ((DEBUG_INFO, "QueryCapsuleCapabilities:%r\n", Status));
    if (EFI_ERROR (Status) || (BufferSize > MaxCapsuleSize)) {
      FreePool (Buffer);
      continue;
    }
    BlockDescriptors = AllocateRuntimeZeroPool (2 * sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR));
    if (BlockDescriptors == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }
    BlockDescriptors->Union.DataBlock    = (UINTN)Buffer;
    BlockDescriptors->Length  = BufferSize;
    if ((CapsuleHeaderArray[0]->Flags & CAPSULE_FLAGS_PERSIST_ACROSS_RESET) != 0) {
      Status = gRT->UpdateCapsule(CapsuleHeaderArray,1,(UINTN) BlockDescriptors);
      if (Status != EFI_SUCCESS) {
        DEBUG ((DEBUG_ERROR, "UpdateCapsule %r\n", Status));
        continue;
      }
      //
      // For capsule who has reset flag, after calling UpdateCapsule service,triger a
      // system reset to process capsule persist across a system reset.
      //
      gRT->ResetSystem (ResetType, EFI_SUCCESS, 0, NULL);
    } else {
      //
      // For capsule who has no reset flag, only call UpdateCapsule Service without a
      // system reset. The service will process the capsule immediately.
      //
      Status = gRT->UpdateCapsule(CapsuleHeaderArray,1,(UINTN) BlockDescriptors);
      if (Status != EFI_SUCCESS) {
        DEBUG ((DEBUG_ERROR, "UpdateCapsule %r\n", Status));
      }
    }
    FreePool (Buffer);
  }

  DEBUG ((DEBUG_INFO, "AutomaticFirmwareUpdateHandler exit\n"));
  return Status;
}

EFI_STATUS
SetMorControl (
  VOID
  )
{
  UINT8                        MorControl;
  UINTN                        VariableSize;
  EFI_STATUS                   Status;

  VariableSize = sizeof (MorControl);
  MorControl = 1;

  Status = gRT->SetVariable(
                  MEMORY_OVERWRITE_REQUEST_VARIABLE_NAME,
                  &gEfiMemoryOverwriteControlDataGuid,
                  EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                  VariableSize,
                  &MorControl
                  );

  return Status;
}

/**

  This routine is called to see if there are any capsules we need to process.
  If the boot mode is not UPDATE, then we do nothing. Otherwise find the
  capsule HOBS and produce firmware volumes for them via the DXE service.
  Then call the dispatcher to dispatch drivers from them. Finally, check
  the status of the updates.

  This function should be called by BDS in case we need to do some
  sort of processing even if there is no capsule to process. We
  need to do this if an earlier update went away and we need to
  clear the capsule variable so on the next reset PEI does not see it and
  think there is a capsule available.

  @retval EFI_INVALID_PARAMETER   boot mode is not correct for an update
  @retval EFI_SUCCESS             There is no error when processing capsule

**/
EFI_STATUS
EFIAPI
EfiPlatformBootManagerProcessCapsules (
  VOID
  )
{
  EFI_STATUS                  Status;
  EFI_PEI_HOB_POINTERS        HobPointer;
  EFI_CAPSULE_HEADER          *CapsuleHeader;
  UINT32                      Size;
  UINT32                      CapsuleNumber;
  UINT32                      CapsuleTotalNumber;
  CRASHDUMP_CAPSULE_TABLE     *CapsuleTable;
  UINT32                      Index;
  UINT32                      CacheIndex;
  UINT32                      CacheNumber;
  VOID                        **CapsulePtr;
  VOID                        **CapsulePtrCache;
  EFI_GUID                    *CapsuleGuidCache;
  EFI_HANDLE                  Handle;
  CAPSULE_STATUS_PROTOCOL     *CapsuleStatusProtocol;
  BOOLEAN                     EcUpdated;
  BOOLEAN                     BaseEcUpdated;
  BOOLEAN                     SystemResetNeeded;
  EFI_STATUS                  *CapsuleStatusArray;
  BOOLEAN                     DisplayCapsuleExist;
  WDT_PROTOCOL                *WdtProtocol;
  PCH_RESET_DATA              ResetData;

  CapsuleNumber     = 0;
  CapsuleTotalNumber = 0;
  CacheIndex        = 0;
  CacheNumber       = 0;
  CapsulePtr        = NULL;
  CapsulePtrCache   = NULL;
  CapsuleGuidCache  = NULL;
  SystemResetNeeded = FALSE;
  EcUpdated = FALSE;
  BaseEcUpdated = FALSE;
  DisplayCapsuleExist = FALSE;
  Status = EFI_SUCCESS;

  //
  // We don't do anything else if the boot mode is not flash-update
  //
  ASSERT (gBootMode == BOOT_ON_FLASH_UPDATE);

  //
  // Find all capsule images from hob
  //
  HobPointer.Raw = GetHobList ();
  while ((HobPointer.Raw = GetNextHob (EFI_HOB_TYPE_UEFI_CAPSULE, HobPointer.Raw)) != NULL) {
    CapsuleTotalNumber ++;
    HobPointer.Raw = GET_NEXT_HOB (HobPointer);
  }

  if (CapsuleTotalNumber == 0) {
    //
    // We didn't find a hob
    //
    return EFI_NOT_FOUND;
  }

  CapsuleStatusProtocol = AllocatePool (sizeof (CAPSULE_STATUS_PROTOCOL));
  if (CapsuleStatusProtocol == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  CapsuleStatusProtocol->CapsuleStatus = EFI_NOT_STARTED;
  CapsuleStatusProtocol->SystemResetNeeded = FALSE;
  CapsuleStatusProtocol->EcUpdated = FALSE;
  CapsuleStatusProtocol->BaseEcUpdated = FALSE;
  Handle = NULL;
  Status = gBS->InstallProtocolInterface (
                  &Handle,
                  &gCapsuleStatusProtocolGuid,
                  EFI_NATIVE_INTERFACE,
                  CapsuleStatusProtocol
                  );
  //
  // Clear the existing Variable "CapsuleXXXX"
  //
  ClearCapsuleStatusVariable ();

  //
  // Init temp Capsule Data table.
  //
  CapsulePtr       = (VOID **) AllocateZeroPool (sizeof (VOID *) * CapsuleTotalNumber);
  ASSERT (CapsulePtr != NULL);
  if (CapsulePtr == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  CapsulePtrCache  = (VOID **) AllocateZeroPool (sizeof (VOID *) * CapsuleTotalNumber);
  ASSERT (CapsulePtrCache != NULL);
  if (CapsulePtrCache == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  CapsuleGuidCache = (EFI_GUID *) AllocateZeroPool (sizeof (EFI_GUID) * CapsuleTotalNumber);
  ASSERT (CapsuleGuidCache != NULL);
  if (CapsuleGuidCache == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  CapsuleStatusArray = (EFI_STATUS *) AllocateZeroPool (sizeof (EFI_STATUS) * CapsuleTotalNumber);
  ASSERT (CapsuleGuidCache != NULL);
  if (CapsuleStatusArray == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  //
  // TPM Shutdown
  //
  Status = Tpm2Shutdown (TPM_SU_CLEAR);
  if (RETURN_ERROR (Status)) {
    DEBUG ((DEBUG_WARN, "%a(%d) Tpm2Shutdown (TPM_SU_CLEAR), Status = %r\n", __FILE__, __LINE__, Status));
  }

  //
  // Find all capsule images from hob
  //
  HobPointer.Raw = GetHobList ();
  while ((HobPointer.Raw = GetNextHob (EFI_HOB_TYPE_UEFI_CAPSULE, HobPointer.Raw)) != NULL) {
    CapsulePtr [CapsuleNumber++] = (VOID *) (UINTN) HobPointer.Capsule->BaseAddress;
    HobPointer.Raw = GET_NEXT_HOB (HobPointer);
  }

  //
  //Check the capsule flags,if contains CAPSULE_FLAGS_POPULATE_SYSTEM_TABLE, install
  //capsuleTable to configure table with EFI_CAPSULE_GUID
  //

  //
  // Capsules who have CAPSULE_FLAGS_POPULATE_SYSTEM_TABLE always are used for operating
  // System to have information persist across a system reset. EFI System Table must
  // point to an array of capsules that contains the same CapsuleGuid value. And agents
  // searching for this type capsule will look in EFI System Table and search for the
  // capsule's Guid and associated pointer to retrieve the data. Two steps below describes
  // how to sorting the capsules by the unique guid and install the array to EFI System Table.
  // Firstly, Loop for all coalesced capsules, record unique CapsuleGuids and cache them in an
  // array for later sorting capsules by CapsuleGuid.
  //
  for (Index = 0; Index < CapsuleTotalNumber; Index++) {
    CapsuleStatusArray [Index] = EFI_UNSUPPORTED;
    CapsuleHeader = (EFI_CAPSULE_HEADER*) CapsulePtr [Index];
    if ((CapsuleHeader->Flags & CAPSULE_FLAGS_POPULATE_SYSTEM_TABLE) != 0) {
      //
      // For each capsule, we compare it with known CapsuleGuid in the CacheArray.
      // If already has the Guid, skip it. Whereas, record it in the CacheArray as
      // an additional one.
      //
      CacheIndex = 0;
      while (CacheIndex < CacheNumber) {
        if (CompareGuid(&CapsuleGuidCache[CacheIndex],&CapsuleHeader->CapsuleGuid)) {
          break;
        }
        CacheIndex++;
      }
      if (CacheIndex == CacheNumber) {
        CopyMem(&CapsuleGuidCache[CacheNumber++],&CapsuleHeader->CapsuleGuid,sizeof(EFI_GUID));
      }
    }
  }

  //
  // Secondly, for each unique CapsuleGuid in CacheArray, gather all coalesced capsules
  // whose guid is the same as it, and malloc memory for an array which preceding
  // with UINT32. The array fills with entry point of capsules that have the same
  // CapsuleGuid, and UINT32 represents the size of the array of capsules. Then install
  // this array into EFI System Table, so that agents searching for this type capsule
  // will look in EFI System Table and search for the capsule's Guid and associated
  // pointer to retrieve the data.
  //
  CacheIndex = 0;
  while (CacheIndex < CacheNumber) {
    CapsuleNumber = 0;
    for (Index = 0; Index < CapsuleTotalNumber; Index++) {
      CapsuleHeader = (EFI_CAPSULE_HEADER*) CapsulePtr [Index];
      if ((CapsuleHeader->Flags & CAPSULE_FLAGS_POPULATE_SYSTEM_TABLE) != 0) {
        if (CompareGuid (&CapsuleGuidCache[CacheIndex], &CapsuleHeader->CapsuleGuid)) {
          //
          // Cache Caspuleheader to the array, this array is uniqued with certain CapsuleGuid.
          //
          CapsulePtrCache[CapsuleNumber++] = (VOID*)CapsuleHeader;
          //
          // When a Capsule is listed in CapsulePtrCache, it will be reported in ConfigurationTable
          // So, report the CapsuleStatus as "processed successfully".
          //
          CapsuleStatusArray [Index] = EFI_SUCCESS;
        }
      }
    }
    if (CapsuleNumber != 0) {
      Size = sizeof(CRASHDUMP_CAPSULE_TABLE) + (CapsuleNumber - 1) * sizeof(VOID*);
      CapsuleTable = AllocateRuntimePool (Size);
      ASSERT (CapsuleTable != NULL);
      if (CapsuleTable == NULL) {
        return EFI_OUT_OF_RESOURCES;
      }
      CapsuleTable->CapsuleArrayNumber =  CapsuleNumber;
      CopyMem(&CapsuleTable->CapsulePtr[0], CapsulePtrCache, CapsuleNumber * sizeof(VOID*));
      Status = gBS->InstallConfigurationTable (&CapsuleGuidCache[CacheIndex], (VOID*)CapsuleTable);
      ASSERT_EFI_ERROR (Status);
    }
    CacheIndex++;
  }


  if (DisplayCapsuleExist == FALSE) {
    //
    // Display Capsule not found. Display the default string.
    //
    gST->ConOut->OutputString(gST->ConOut, L"Updating the firmware do not un-plug!\r\n");
  }


  for (Index = 0; Index < CapsuleTotalNumber; Index++) {
    Status = CreateCapsuleStatusVariable (Index, (EFI_CAPSULE_HEADER *) CapsulePtr [Index], CapsuleStatusArray [Index]);
  }
  //
  // Free the allocated temp memory space.
  //
  FreePool (CapsuleGuidCache);
  FreePool (CapsulePtrCache);
  FreePool (CapsulePtr);

  //
  // Set MOR control bit to clear memory before initiating system reset, so that memory will be cleared in next boot.
  //
  SetMorControl();

  if (mPlatformInfo->EcPresent) {
    if (EcUpdated) {
      DEBUG ((DEBUG_INFO, "Send EC Reset Cmd!\n"));
      Status = EcReset (); // <-0x0B
      gBS->Stall (2000000); // 2sec
      DEBUG ((DEBUG_INFO, "Pending Reset after 2Sec EC Reset!\n"));
    } else if(BaseEcUpdated) {
      DEBUG ((DEBUG_INFO, "Reseting EC via Reset Seq \n"));
      Status = EcResetEcInNormalMode();
    }
  }

  if (SystemResetNeeded) {
    DEBUG ((DEBUG_INFO, "Initiating System Reset \n"));
    Status = gBS->LocateProtocol (&gWdtProtocolGuid, NULL, (VOID **) &WdtProtocol);
    if (!EFI_ERROR (Status)) {
      WdtProtocol->AllowKnownReset();
    }
    DEBUG ((DEBUG_INFO, "Global Reset requested by FW EOP ACK %r\n", Status));
    CopyMem (&ResetData.Guid, &gPchGlobalResetGuid, sizeof (EFI_GUID));
    StrCpyS (ResetData.Description, PCH_RESET_DATA_STRING_MAX_LENGTH, PCH_PLATFORM_SPECIFIC_RESET_STRING);
    gRT->ResetSystem (EfiResetPlatformSpecific, EFI_SUCCESS, sizeof (PCH_RESET_DATA), &ResetData);
  }

  return Status;
}

